package com.onefish.toolkit.cwe;

import org.apache.commons.exec.*;

import java.io.ByteArrayOutputStream;
import java.util.Arrays;

public class SecureCommandExecutor {
    private static final String[] blackList = new String[]{"&", "&&", "|", "||", ";", "(", ")"};

    public static void main(String[] args) throws Exception {
        executeSecureCommand(true, "echo", "hello", "world", "hhhhhhhhhhhhhhhhhhhhhhhhhhhhh");
    }

    public static void executeSecureCommand(boolean specialCharCheck, String... command) {
        try {
            CommandLine cmdLine = null;
            for (int i = 0; i < command.length; i++) {
                int temp = i;
                // 开启命令行黑名单校验
                if (specialCharCheck) {
                    Arrays.stream(blackList).forEach(black -> {
                        if (command[temp].contains(black)) {
                            throw new IllegalArgumentException("执行命令中包含非法字符: " + black);
                        }
                    });
                }
                //  获取第一个命令参数初始化操作
                if (cmdLine == null){
                    cmdLine = CommandLine.parse(command[temp]);
                    continue;
                }
                cmdLine.addArgument(command[temp]);
            }
            Executor executor = new DefaultExecutor();
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            PumpStreamHandler streamHandler = new PumpStreamHandler(outputStream);
            executor.setStreamHandler(streamHandler);
            CustomResultHandler resultHandler = new CustomResultHandler(outputStream);
            executor.execute(cmdLine, resultHandler);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
